from sqlalchemy import MetaData
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase

from rosa.scarborough_fair import settings


class Base(DeclarativeBase):
    metadata = MetaData()


class SqlHelper:
    engine = None
    session: async_sessionmaker | None = None
    model_base = Base

    def __init__(self):
        # 如果需要数据库的项目，则初始化orm
        if settings.DB_ENABLE:
            self.init_orm()

    def init_orm(self) -> None:
        db_url = f"{settings.DB_ENGINE}://{settings.DB_USER}:{settings.DB_PASSWORD}@{settings.DB_HOST}:{settings.DB_PORT}/{settings.DB_DATABASE}"
        self.engine = create_async_engine(db_url, future=True, echo=settings.DB_ECHO, pool_pre_ping=True, pool_recycle=3600)
        self.session = async_sessionmaker(self.engine, expire_on_commit=False)

    def get_session(self) -> async_sessionmaker:
        if not self.session:
            raise ValueError("未初始化数据库连接，请检查配置文件是否开启数据库功能")
        return self.session


sql_helper = SqlHelper()

ModelBase = sql_helper.model_base
